package com.codeanalysissys.backend.python;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Map;

import com.codeanalysissys.backend.graphs.ast.ASNode;
import com.codeanalysissys.backend.graphs.ast.AbstractSyntaxTree;
import com.codeanalysissys.backend.java.parser.JavaLexer;
import com.codeanalysissys.backend.java.parser.JavaParser;
import com.codeanalysissys.backend.python.parser.PythonLexer;
import com.codeanalysissys.backend.python.parser.PythonParser;
import com.codeanalysissys.backend.python.parser.PythonParserBaseVisitor;
import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.misc.Interval;
import org.antlr.v4.runtime.tree.ParseTree;
import ghaffarian.nanologger.Logger;

import java.util.LinkedHashMap;
import org.antlr.v4.runtime.tree.TerminalNode;

/**
 * Abstract Syntax Tree (AST) builder for Python programs.
 * A Python parser generated via ANTLRv4 is used for this purpose.
 * This implementation is based on ANTLRv4's Visitor pattern.
 *
 * @author WenLiang Peng
 */
public class PythonASTBuilder {

    /**
     * ‌Build and return the Abstract Syntax Tree (AST) for the given Python source file.
     */
    public static AbstractSyntaxTree build(String pythonFile) throws IOException {
        return build(new File(pythonFile));
    }

    public static AbstractSyntaxTree build(InputStream inFile) throws IOException {
        ANTLRInputStream input = new ANTLRInputStream(inFile);
        PythonLexer lexer = new PythonLexer(input);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        PythonParser parser = new PythonParser(tokens);
        ParseTree tree = parser.root();
        return build(tree, null, null);
    }

    /**
     * ‌Build and return the Abstract Syntax Tree (AST) for the given Python source file.
     */
    public static AbstractSyntaxTree build(File pythonFile) throws IOException {
        if (!pythonFile.getName().endsWith(".py"))
            throw new IOException("Not a Python File!");
        InputStream inFile = new FileInputStream(pythonFile);
        ANTLRInputStream input = new ANTLRInputStream(inFile);                   //将源代码文件输入流转换为Antlr的输入流
        PythonLexer lexer = new PythonLexer(input);                                  //Antlr生成的词法分析器
        CommonTokenStream tokens = new CommonTokenStream(lexer);                 //将词法分析结果转换为token流
        PythonParser parser = new PythonParser(tokens);                              //Antlr生成的语法分析器
        ParseTree tree = parser.root();                                          //Python语法规则中的根规则-解析器已生成
        return build(pythonFile.getPath(), tree, null, null);
    }

    public static AbstractSyntaxTree build(ParseTree tree, String propKey, Map<ParserRuleContext, Object> ctxProps) {
        PythonASTBuilder.AbstractSyntaxVisitor visitor = new PythonASTBuilder.AbstractSyntaxVisitor(propKey, ctxProps);         //采用Antlr的visitor遍历
        return visitor.build(tree);
    }

    /**
     * ‌Build and return the Abstract Syntax Tree (AST) for the given Parse-Tree.
     * The 'ctxProps' map includes contextual-properties for particular nodes
     * in the parse-tree, which can be used for linking this graph with other
     * graphs by using the same parse-tree and the same contextual-properties.
     */
    public static AbstractSyntaxTree build(String filePath, ParseTree tree,
                                           String propKey, Map<ParserRuleContext, Object> ctxProps) {
        PythonASTBuilder.AbstractSyntaxVisitor visitor = new PythonASTBuilder.AbstractSyntaxVisitor(filePath, propKey, ctxProps);         //采用Antlr的visitor遍历
        Logger.debug("Visitor building AST of: " + filePath);
        return visitor.build(tree);
    }

    /**
     * Visitor class which constructs the AST for a given ParseTree.
     */
    private static class AbstractSyntaxVisitor extends PythonParserBaseVisitor<String> {

        private String propKey;
        private int importNum = 0;
        private String typeModifier;
        private String memberModifier;
        private Deque<ASNode> parentStack;
        private final AbstractSyntaxTree AST;
        private Map<String, String> vars, fields, methods;
        private int varsCounter, fieldsCounter, methodsCounter;
        private Map<ParserRuleContext, Object> contexutalProperties;

        public AbstractSyntaxVisitor(String filePath, String propKey, Map<ParserRuleContext, Object> ctxProps) {
            parentStack = new ArrayDeque<>();
            AST = new AbstractSyntaxTree(filePath);
            this.propKey = propKey;
            contexutalProperties = ctxProps;
            vars = new LinkedHashMap<>();
            fields = new LinkedHashMap<>();
            methods = new LinkedHashMap<>();
            varsCounter = 0;
            fieldsCounter = 0;
            methodsCounter = 0;
        }

        public AbstractSyntaxVisitor(String propKey, Map<ParserRuleContext, Object> ctxProps) {
            parentStack = new ArrayDeque<>();
            AST = new AbstractSyntaxTree();
            this.propKey = propKey;
            contexutalProperties = ctxProps;
            vars = new LinkedHashMap<>();
            fields = new LinkedHashMap<>();
            methods = new LinkedHashMap<>();
            varsCounter = 0;
            fieldsCounter = 0;
            methodsCounter = 0;
        }

        public AbstractSyntaxTree build(ParseTree tree) {
            // : (single_input
            //    | file_input
            //    | eval_input)? EOF
            // deal with the root node first
            PythonParser.RootContext rootCntx = (PythonParser.RootContext) tree;
            AST.root.setCode(AST.filePath != null ? new File(AST.filePath).getName() : "root");
            parentStack.push(AST.root);                              // store the parent node,now is the ast root node
            // no package and import and so on in python
            if (rootCntx.file_input() != null) {
                visit(rootCntx.file_input());
            }
            if (rootCntx.single_input() != null) {
                visit(rootCntx.single_input());
            }
            if (rootCntx.eval_input() != null) {
                visit(rootCntx.eval_input());
            }
            parentStack.pop();
            vars.clear();
            fields.clear();
            methods.clear();
            importNum = 0;
            return AST;
        }

        //=====================================================================//
        //                           DECLARATIONS                              //
        //=====================================================================//

        @Override
        public String visitImport_stmt(PythonParser.Import_stmtContext ctx) {
            // IMPORT dotted_as_names
            if (++importNum == 1) {
                // the first to create field node
                ASNode imports = new ASNode(ASNode.Type.IMPORTS);
                imports.setLineOfCode(ctx.getStart().getLine());
                AST.addVertex(imports);
                AST.addEdge(AST.root, imports);
                parentStack.push(imports);
            }
            // now deal with the IMPORT Node
            ASNode node = new ASNode(ASNode.Type.IMPORT);
            String code = "import ";
            for (int i = 0; i < ctx.dotted_as_names().COMMA().size(); i++){
                if (i == 0){
                    code = code + ctx.dotted_as_names().dotted_as_name(i).getText();
                }else {
                    code = code + ", " + ctx.dotted_as_names().dotted_as_name(i).getText();
                }
            }
            node.setCode(code);
            node.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(node);
            AST.addEdge(parentStack.peek(), node);
            return "";
        }

        @Override
        public String visitFrom_stmt(PythonParser.From_stmtContext ctx) {
            // FROM ((DOT | ELLIPSIS)* dotted_name | (DOT | ELLIPSIS)+)
            //      IMPORT (STAR | OPEN_PAREN import_as_names CLOSE_PAREN | import_as_names)
            if (++importNum == 1) {
                // the first to create field node
                ASNode imports = new ASNode(ASNode.Type.IMPORTS);
                imports.setLineOfCode(ctx.getStart().getLine());
                AST.addVertex(imports);
                AST.addEdge(AST.root, imports);
                parentStack.push(imports);
            }
            // now deal with the IMPORT Node
            ASNode node = new ASNode(ASNode.Type.IMPORT);
            String code = "from ";
            for(int i = 0; i < ctx.DOT().size(); i++){
                code = code + ctx.DOT(i).getText();
            }
            code = (ctx.dotted_name() != null) ? code + ctx.dotted_name().getText() : code;
            code = code + " import ";
            if (ctx.STAR() != null){
                code = code + ctx.STAR().getText();
            }else{
                if (ctx.import_as_names() != null){
                    assert ctx.import_as_names() instanceof PythonParser.Import_as_namesContext;
                    code = (ctx.OPEN_PAREN() != null) ? code + '(' : code;
                    int com_length = ctx.import_as_names().COMMA().size();
                    for (int i = 0; i < ctx.import_as_names().import_as_name().size(); i++){
                        PythonParser.Import_as_nameContext temAsName = ctx.import_as_names().import_as_name(i);
                        code = (temAsName.AS() != null) ? code + temAsName.name(0).getText() + " as " + temAsName.name(1).getText() : code;
                        if (com_length-- != 0) code = code + ',';
                    }
                }else{
                    //这种情况不存在
                }
            }
            node.setCode(code);
            node.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(node);
            AST.addEdge(parentStack.peek(), node);
            return "";
        }

        @Override
        public String visitClassdef(PythonParser.ClassdefContext ctx) {
            // CLASS name (OPEN_PAREN arglist? CLOSE_PAREN)? COLON suite
            ASNode classNode = new ASNode(ASNode.Type.CLASS);
            classNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(classNode);
            // judge whether the IMPORTS node exist
            if (parentStack.peek().getType().toString().equals("IMPORTS"))
                parentStack.pop();
            if (parentStack.peek().getType().toString().equals("CLASS"))
                parentStack.pop();
            AST.addEdge(parentStack.peek(), classNode);
            // now deal with the AST node
            ASNode modifierNode = new ASNode(ASNode.Type.MODIFIER);
            modifierNode.setCode(typeModifier);                            // the class type is null
            modifierNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(modifierNode);
            AST.addEdge(classNode, modifierNode);
            // now add the Name tag
            ASNode nameNode = new ASNode(ASNode.Type.NAME);
            String className = ctx.name().getText();
            nameNode.setCode(className);
            nameNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(nameNode);
            AST.addEdge(classNode, nameNode);
            // now add the EXTENTS field
            if (ctx.arglist() != null) {
                // create the field node first
                ASNode extendsNode = new ASNode(ASNode.Type.EXTENDS);
                extendsNode.setCode("EXTENDS");
                extendsNode.setLineOfCode(ctx.arglist().getStart().getLine());
                AST.addVertex(extendsNode);
                AST.addEdge(classNode, extendsNode);
                // add the edge from the extents node to extent node
                for (int i = 0; i < ctx.arglist().argument().size(); i++) {
                    ASNode extendNode = new ASNode(ASNode.Type.NAME);
                    extendNode.setCode(ctx.arglist().argument(i).getText());
                    extendNode.setLineOfCode(ctx.arglist().argument(i).getStart().getLine());
                    AST.addVertex(extendNode);
                    AST.addEdge(extendsNode, extendNode);
                }
            }
            // and then deal with the body
            parentStack.push(classNode);
            visit(ctx.suite());
            parentStack.pop();
            return "";
        }

        @Override
        public String visitFuncdef(PythonParser.FuncdefContext ctx) {
            // ASYNC? DEF name OPEN_PAREN typedargslist? CLOSE_PAREN (ARROW test)? COLON suite
            // create the FUNCTION node
            ASNode functionNode = new ASNode(ASNode.Type.FUNCTION);
            functionNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(functionNode);
            AST.addEdge(parentStack.peek(), functionNode);
            parentStack.push(functionNode);
            // now deal with the tag
            ASNode modifierNode = new ASNode(ASNode.Type.MODIFIER);
            memberModifier = (ctx.ASYNC() != null) ? "async" : null;
            modifierNode.setCode(memberModifier);
            modifierNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(modifierNode);
            AST.addEdge(parentStack.peek(), modifierNode);
            // and then deal with the return node
            ASNode retNode = new ASNode(ASNode.Type.RETURN);
            retNode.setCode(ctx.test() != null ? ctx.test().getText() : null);
            retNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(retNode);
            AST.addEdge(parentStack.peek(), retNode);
            // and then deal with the name information
            ++methodsCounter;
            ASNode nameNode = new ASNode(ASNode.Type.NAME);
            String methodName = ctx.name().getText();
            // store the method imformation
            String normalized = "$METHOD_" + methodsCounter;
            methods.put(methodName, normalized);
            nameNode.setCode(methodName);
            nameNode.setNormalizedCode(normalized);
            nameNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(nameNode);
            AST.addEdge(parentStack.peek(), nameNode);
            // now deal with arglist to get params
            if (ctx.typedargslist() != null) {
                // create the params field first
                ASNode paramsNode = new ASNode(ASNode.Type.PARAMS);
                paramsNode.setLineOfCode(ctx.typedargslist().getStart().getLine());
                AST.addVertex(paramsNode);
                AST.addEdge(parentStack.peek(), paramsNode);
                parentStack.push(paramsNode);
                // first deal with the def_parameters
                for (PythonParser.Def_parametersContext temDefs : ctx.typedargslist().def_parameters()) {
                    for (PythonParser.Def_parameterContext temDef : temDefs.def_parameter()) {
                        // create the var node
                        ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                        varNode.setLineOfCode(temDef.getStart().getLine());
                        AST.addVertex(varNode);
                        AST.addEdge(parentStack.peek(), varNode);
                        // now create the name node
                        ++varsCounter;
                        ASNode name = new ASNode(ASNode.Type.NAME);
                        normalized = "$VARL_" + varsCounter;
                        vars.put(temDef.getText(), normalized);
                        name.setCode(temDef.getText());
                        name.setNormalizedCode(normalized);
                        name.setLineOfCode(temDef.getStart().getLine());
                        AST.addVertex(name);
                        AST.addEdge(varNode, name);
                    }
                }
                // and then deal with the arg and karg
                if (ctx.typedargslist().args() != null) {
                    // create the var node
                    ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                    varNode.setLineOfCode(ctx.typedargslist().args().getStart().getLine());
                    AST.addVertex(varNode);
                    AST.addEdge(parentStack.peek(), varNode);
                    // now create the name node
                    ++varsCounter;
                    ASNode name = new ASNode(ASNode.Type.NAME);
                    normalized = "$VARL_" + varsCounter;
                    vars.put(ctx.typedargslist().args().getText(), normalized);
                    name.setCode(ctx.typedargslist().args().getText());
                    name.setNormalizedCode(normalized);
                    name.setLineOfCode(ctx.typedargslist().args().getStart().getLine());
                    AST.addVertex(name);
                    AST.addEdge(varNode, name);
                }
                if (ctx.typedargslist().kwargs() != null) {
                    // create the var node
                    ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                    varNode.setLineOfCode(ctx.typedargslist().kwargs().getStart().getLine());
                    AST.addVertex(varNode);
                    AST.addEdge(parentStack.peek(), varNode);
                    // now create the name node
                    ++varsCounter;
                    ASNode name = new ASNode(ASNode.Type.NAME);
                    normalized = "$VARL_" + varsCounter;
                    vars.put(ctx.typedargslist().kwargs().getText(), normalized);
                    name.setCode(ctx.typedargslist().kwargs().getText());
                    name.setNormalizedCode(normalized);
                    name.setLineOfCode(ctx.typedargslist().kwargs().getStart().getLine());
                    AST.addVertex(name);
                    AST.addEdge(varNode, name);
                }
                parentStack.pop();
            }
            // now deal with the body
            if (ctx.suite() != null) {
                ASNode suiteBody = new ASNode(ASNode.Type.SUITE);
                suiteBody.setLineOfCode(ctx.suite().getStart().getLine());
                AST.addVertex(suiteBody);
                AST.addEdge(parentStack.peek(), suiteBody);
                parentStack.push(suiteBody);
                visitChildren(ctx.suite());
                parentStack.pop();
                resetLocalVars();
            }
            return "";
        }

        @Override
        public String visitExpr_stmt(PythonParser.Expr_stmtContext ctx) {
            // : testlist_star_expr assign_part?
            if (ctx.assign_part() != null) {
                if (ctx.assign_part().ASSIGN().size() == 1) {
                    // if just one "=" then create a VAR node first
                    ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                    varNode.setLineOfCode(ctx.getStart().getLine());
                    AST.addVertex(varNode);
                    AST.addEdge(parentStack.peek(), varNode);
                    //
                    ++varsCounter;
                    ASNode nameNode = new ASNode(ASNode.Type.NAME);
                    String normalized = "$VARL_" + varsCounter;
                    vars.put(ctx.testlist_star_expr().getText(), normalized);
                    nameNode.setCode(ctx.testlist_star_expr().getText());
                    nameNode.setNormalizedCode(normalized);
                    nameNode.setLineOfCode(ctx.testlist_star_expr().getStart().getLine());
                    AST.addVertex(nameNode);
                    AST.addEdge(varNode, nameNode);
                    //
                    ASNode initNode = new ASNode(ASNode.Type.INIT_VALUE);
                    initNode.setCode("= " + getOriginalCodeText(ctx.assign_part().testlist_star_expr(0)));
                    initNode.setNormalizedCode("= $VAR");
                    initNode.setLineOfCode(ctx.getStart().getLine());
                    AST.addVertex(initNode);
                    AST.addEdge(varNode, initNode);
                } else if (ctx.assign_part().ASSIGN().size() > 1) {
                    int last = ctx.assign_part().testlist_star_expr().size() - 1;
                    // exist many assign part
                    for (int i = 0; i < ctx.assign_part().testlist_star_expr().size(); i++) {
                        if (i != ctx.assign_part().testlist_star_expr().size() - 1) {
                            ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                            varNode.setLineOfCode(ctx.getStart().getLine());
                            AST.addVertex(varNode);
                            AST.addEdge(parentStack.peek(), varNode);
                            //
                            ++varsCounter;
                            ASNode nameNode = new ASNode(ASNode.Type.NAME);
                            String normalized = "$VARL_" + varsCounter;
                            vars.put(ctx.assign_part().testlist_star_expr(i).getText(), normalized);
                            nameNode.setCode(ctx.assign_part().testlist_star_expr(i).getText());
                            nameNode.setNormalizedCode(normalized);
                            nameNode.setLineOfCode(ctx.assign_part().testlist_star_expr(i).getStart().getLine());
                            AST.addVertex(nameNode);
                            AST.addEdge(varNode, nameNode);
                            //
                            ASNode initNode = new ASNode(ASNode.Type.INIT_VALUE);
                            initNode.setCode("= " + getOriginalCodeText(ctx.assign_part().testlist_star_expr(last)));
                            initNode.setNormalizedCode("= $VAR");
                            initNode.setLineOfCode(ctx.assign_part().testlist_star_expr(i).getStart().getLine());
                            AST.addVertex(initNode);
                            AST.addEdge(varNode, initNode);
                        }
                        for (PythonParser.TestContext temTest : ctx.testlist_star_expr().test()) {
                            ASNode varNode = new ASNode(ASNode.Type.VARIABLE);
                            varNode.setLineOfCode(ctx.getStart().getLine());
                            AST.addVertex(varNode);
                            AST.addEdge(parentStack.peek(), varNode);
                            //
                            ++varsCounter;
                            ASNode nameNode = new ASNode(ASNode.Type.NAME);
                            String normalized = "$VARL_" + varsCounter;
                            vars.put(temTest.getText(), normalized);
                            nameNode.setCode(temTest.getText());
                            nameNode.setNormalizedCode(normalized);
                            nameNode.setLineOfCode(temTest.getStart().getLine());
                            AST.addVertex(nameNode);
                            AST.addEdge(varNode, nameNode);
                            //
                            ASNode initNode = new ASNode(ASNode.Type.INIT_VALUE);
                            initNode.setCode("= " + getOriginalCodeText(ctx.assign_part().testlist_star_expr(last)));
                            initNode.setNormalizedCode("= $VAR");
                            initNode.setLineOfCode(ctx.assign_part().testlist_star_expr(i).getStart().getLine());
                            AST.addVertex(initNode);
                            AST.addEdge(varNode, initNode);
                        }
                    }
                }
            } else {
                // deal with like the normal statement
                visitStatement(ctx, "normal statement");
            }
            return "";
        }

        @Override
        public String visitTry_stmt(PythonParser.Try_stmtContext ctx) {
            // TRY COLON suite (except_clause+ else_clause? finally_clause? | finally_clause)
            // create a try node first
            ASNode tryNode = new ASNode(ASNode.Type.TRY);
            tryNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(tryNode);
            AST.addEdge(parentStack.peek(), tryNode);                   // create a edge from method node to try node
            // now deal with the block
            ASNode trySuite = new ASNode(ASNode.Type.SUITE);
            trySuite.setLineOfCode(ctx.suite().getStart().getLine());
            AST.addVertex(trySuite);
            AST.addEdge(tryNode, trySuite);
            parentStack.push(trySuite);
            visit(ctx.suite());
            parentStack.pop();
            // EXCEPT (test ({this.CheckVersion(2)}? COMMA name {this.SetVersion(2);} | {this.CheckVersion(3)}? AS name {this.SetVersion(3);})?)? COLON suite
            if (ctx.except_clause() != null && ctx.except_clause().size() > 0) {
                for (PythonParser.Except_clauseContext exceptCtx : ctx.except_clause()) {
                    ASNode exceptNode = new ASNode(ASNode.Type.EXCEPT);
                    exceptNode.setLineOfCode(exceptCtx.getStart().getLine());
                    AST.addVertex(exceptNode);
                    AST.addEdge(tryNode, exceptNode);
                    // deal with the except nodes
                    ++varsCounter;
                    ASNode exceptName = new ASNode(ASNode.Type.NAME);
                    String normalized = "$VARL_" + varsCounter;
                    if (exceptCtx != null && exceptCtx.name() != null) {
                        vars.put(exceptCtx.name().getText(), normalized);
                        exceptName.setCode(exceptCtx.name().getText());
                        exceptName.setNormalizedCode(normalized);
                        exceptName.setLineOfCode(exceptCtx.getStart().getLine());
                    } else {
                        vars.put(exceptCtx.getText(), normalized);
                        exceptName.setCode(exceptCtx.getText());
                        exceptName.setNormalizedCode(normalized);
                        exceptName.setLineOfCode(exceptCtx.getStart().getLine());
                    }
                    AST.addVertex(exceptName);
                    AST.addEdge(exceptNode, exceptName);
                    // now deal with the except block
                    ASNode exceptSuite = new ASNode(ASNode.Type.SUITE);
                    exceptSuite.setLineOfCode(exceptCtx.getStart().getLine());
                    AST.addVertex(exceptSuite);
                    AST.addEdge(exceptNode, exceptSuite);
                    parentStack.push(exceptSuite);
                    visit(exceptCtx.suite());
                    parentStack.pop();
                }
            }
            // :FINALLY COLON suite
            if (ctx.finally_clause() != null) {
                ASNode finallyNode = new ASNode(ASNode.Type.FINALLY);
                finallyNode.setLineOfCode(ctx.finally_clause().getStart().getLine());
                AST.addVertex(finallyNode);
                AST.addEdge(tryNode, finallyNode);
                parentStack.push(finallyNode);
                visit(ctx.finally_clause().suite());
                parentStack.pop();
            }
            // deal with the else suite(python specification)
            if (ctx.else_clause() != null) {
                ASNode elseNode = new ASNode(ASNode.Type.ELSE);
                elseNode.setLineOfCode(ctx.else_clause().getStart().getLine());
                AST.addVertex(elseNode);
                AST.addEdge(tryNode, elseNode);
                parentStack.push(elseNode);
                visit(ctx.else_clause().suite());
                parentStack.pop();
            }
            return "";
        }

        @Override
        public String visitWith_stmt(PythonParser.With_stmtContext ctx) {
            // ASYNC? WITH with_item (COMMA with_item)* COLON suite
            ASNode withNode = new ASNode(ASNode.Type.WITH);
            withNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(withNode);
            AST.addEdge(parentStack.peek(), withNode);
            // deal with the modifierNode
            ASNode modifierNode = new ASNode(ASNode.Type.MODIFIER);
            memberModifier = (ctx.ASYNC() != null) ? "async" : null;
            modifierNode.setCode(memberModifier);
            modifierNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(modifierNode);
            AST.addEdge(withNode, modifierNode);
            // deal with the cond node
            for (PythonParser.With_itemContext temWith : ctx.with_item()) {
                // test (AS expr)?
                ASNode cond = new ASNode(ASNode.Type.CONDITION);
                cond.setCode(getOriginalCodeText(temWith.test()));
                cond.setNormalizedCode(visit(temWith.test()));
                cond.setLineOfCode(temWith.test().getStart().getLine());
                AST.addVertex(cond);
                AST.addEdge(withNode, cond);
                if (temWith.AS() != null && temWith.expr() != null) {
                    ASNode testName = new ASNode(ASNode.Type.NAME);
                    String code = "as ";
                    code = code + temWith.expr().getText();
                    testName.setCode(code);
                    testName.setNormalizedCode(code);
                    testName.setLineOfCode(temWith.expr().getStart().getLine());
                    AST.addVertex(testName);
                    AST.addEdge(cond, testName);
                }
            }
            // now deal with the suite
            parentStack.push(withNode);
            visit(ctx.suite());
            parentStack.pop();
            return "";
        }

        @Override
        public String visitPrint_stmt(PythonParser.Print_stmtContext ctx) {
            //  {this.CheckVersion(2)}? PRINT ((test (COMMA test)* COMMA?)
            //                       | RIGHT_SHIFT test ((COMMA test)+ COMMA?)) {this.SetVersion(2);}
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitDel_stmt(PythonParser.Del_stmtContext ctx) {
            // DEL exprlist
            visitStatement(ctx, "Del" + visit(ctx.exprlist()));
            return "";
        }

        @Override
        public String visitReturn_stmt(PythonParser.Return_stmtContext ctx) {
            // RETURN testlist?
            if (ctx.testlist() != null)
                visitStatement(ctx, "return " + visit(ctx.testlist()));
            return "RETURN";
        }

        @Override
        public String visitRaise_stmt(PythonParser.Raise_stmtContext ctx) {
            // RAISE (test (COMMA test (COMMA test)?)?)? (FROM test)?
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitYield_stmt(PythonParser.Yield_stmtContext ctx) {
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitGlobal_stmt(PythonParser.Global_stmtContext ctx) {
            // GLOBAL name (COMMA name)*
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitAssert_stmt(PythonParser.Assert_stmtContext ctx) {
            // ASSERT test (COMMA test)?
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitNonlocal_stmt(PythonParser.Nonlocal_stmtContext ctx) {
            // {this.CheckVersion(3)}? NONLOCAL name (COMMA name)* {this.SetVersion(3);
            visitStatement(ctx, ctx.getText());
            return "";
        }

        @Override
        public String visitIf_stmt(PythonParser.If_stmtContext ctx) {
            // //IF cond=test COLON suite elif_clause* else_clause?
            ASNode ifNode = new ASNode(ASNode.Type.IF);
            ifNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(ifNode);
            AST.addEdge(parentStack.peek(), ifNode);
            // and then deal with the cond node
            ASNode cond = new ASNode(ASNode.Type.CONDITION);
            cond.setCode(getOriginalCodeText(ctx.test()));
            cond.setNormalizedCode(ctx.test().getText());
            cond.setLineOfCode(ctx.test().getStart().getLine());
            AST.addVertex(cond);
            AST.addEdge(ifNode, cond);
            // and then deal with the true situation, namely the body
            ASNode thenNode = new ASNode(ASNode.Type.THEN);
            thenNode.setLineOfCode(ctx.suite().getStart().getLine());
            AST.addVertex(thenNode);
            AST.addEdge(ifNode, thenNode);
            parentStack.push(thenNode);
            visit(ctx.suite());
            parentStack.pop();
            // ELIF test COLON suite
            if (ctx.elif_clause() != null && ctx.elif_clause().size() > 0) {
                for (PythonParser.Elif_clauseContext temCtx : ctx.elif_clause()) {
                    ASNode elifNode = new ASNode(ASNode.Type.ELIF);
                    elifNode.setLineOfCode(temCtx.getStart().getLine());
                    AST.addVertex(elifNode);
                    AST.addEdge(ifNode, elifNode);
                    parentStack.push(elifNode);
                    visit(temCtx.suite());
                    parentStack.pop();
                }
            }
            // ELSE COLON suite
            if (ctx.else_clause() != null) {
                ASNode elseNode = new ASNode(ASNode.Type.ELSE);
                elseNode.setLineOfCode(ctx.else_clause().getStart().getLine());
                AST.addVertex(elseNode);
                AST.addEdge(ifNode, elseNode);
                parentStack.push(elseNode);
                visit(ctx.else_clause().suite());
                parentStack.pop();
            }
            return "";
        }
        @Override
        public String visitWhile_stmt(PythonParser.While_stmtContext ctx) {
            // WHILE test COLON suite else_clause?
            ASNode whileNode = new ASNode(ASNode.Type.WHILE);
            whileNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(whileNode);
            AST.addEdge(parentStack.peek(), whileNode);
            // create the condition node
            ASNode cond = new ASNode(ASNode.Type.CONDITION);
            cond.setCode(getOriginalCodeText(ctx.test()));
            cond.setNormalizedCode(visit(ctx.test()));
            cond.setLineOfCode(ctx.test().getStart().getLine());
            AST.addVertex(cond);
            AST.addEdge(whileNode, cond);
            // and then create the suite node and deal with the block
            ASNode suite = new ASNode(ASNode.Type.SUITE);
            suite.setLineOfCode(ctx.suite().getStart().getLine());
            AST.addVertex(suite);
            AST.addEdge(whileNode, suite);
            parentStack.push(suite);
            visit(ctx.suite());
            parentStack.pop();
            // deal with the else block
            if (ctx.else_clause() != null) {
                ASNode elseNode = new ASNode(ASNode.Type.ELSE);
                elseNode.setLineOfCode(ctx.else_clause().getStart().getLine());
                AST.addVertex(elseNode);
                AST.addEdge(whileNode, elseNode);
                parentStack.push(elseNode);
                visit(ctx.else_clause().suite());
                parentStack.pop();
            }
            return "";
        }
        @Override
        public String visitFor_stmt(PythonParser.For_stmtContext ctx) {
            // ASYNC? FOR exprlist IN testlist COLON suite else_clause?
            ASNode forNode = new ASNode(ASNode.Type.FOR_EACH);
            forNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(forNode);
            AST.addEdge(parentStack.peek(), forNode);
            //
            ASNode modifierNode = new ASNode(ASNode.Type.MODIFIER);
            memberModifier = (ctx.ASYNC() != null) ? "async" : null;
            modifierNode.setCode(memberModifier);
            modifierNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(modifierNode);
            AST.addEdge(forNode, modifierNode);
            // set the index name
            ASNode indexName = new ASNode(ASNode.Type.INDEXNAME);
            indexName.setCode(ctx.exprlist().getText());
            indexName.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(indexName);
            AST.addEdge(forNode, indexName);
            // set the original index
            ASNode rawIndex = new ASNode(ASNode.Type.RAWINDEX);
            rawIndex.setCode(ctx.testlist().getText());
            rawIndex.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(rawIndex);
            AST.addEdge(forNode, rawIndex);
            // and then deal with the suite
            ASNode suite = new ASNode(ASNode.Type.SUITE);
            suite.setLineOfCode(ctx.suite().getStart().getLine());
            AST.addVertex(suite);
            AST.addEdge(forNode, suite);
            parentStack.push(suite);
            visit(ctx.suite());
            parentStack.pop();
            // and then deal with the else block
            if (ctx.else_clause() != null) {
                ASNode elseNode = new ASNode(ASNode.Type.ELSE);
                elseNode.setLineOfCode(ctx.else_clause().getStart().getLine());
                AST.addVertex(elseNode);
                AST.addEdge(forNode, elseNode);
                parentStack.push(elseNode);
                visit(ctx.else_clause().suite());
                parentStack.pop();
            }
            return "";
        }
        // cfg statements end
        @Override
        public String visitLogical_test(PythonParser.Logical_testContext ctx) {
//			logical_test
//			       : comparison
//					| NOT logical_test
//					| logical_test op=AND logical_test
//					| logical_test op=OR logical_test
//			;
            if (ctx.comparison() != null) {
                visit(ctx.comparison());
            } else if (ctx.NOT() != null && ctx.logical_test() != null) {
                String temStr = visit(ctx.logical_test(0));
                return "not " + temStr;
            } else if (ctx.op != null && ctx.op.getText().equals("AND")){
                String temStr1 = visit(ctx.logical_test(0));
                String temStr2 = visit(ctx.logical_test(1));
                return temStr1 + " and " + temStr2;
            } else if (ctx.op != null && ctx.op.getText().equals("OR")){
                String temStr1 = visit(ctx.logical_test(0));
                String temStr2 = visit(ctx.logical_test(1));
                return temStr1 + " or " + temStr2;
            }
            return "";
        }

        @Override
        public String visitComparison(PythonParser.ComparisonContext ctx) {
            //comparison
            //    : comparison (LESS_THAN | GREATER_THAN | EQUALS | GT_EQ | LT_EQ | NOT_EQ_1 | NOT_EQ_2 | optional=NOT? IN | IS optional=NOT?) comparison
            //    | expr
            //    ;
            if (ctx.expr() != null) {
                String temExpr = visit(ctx.expr());
                return temExpr;
            } else if (ctx.comparison() != null) {
                // situation two
                // situation one
                String temStr1 = visit(ctx.comparison(0));
                String temStr2 = visit(ctx.comparison(1));
                if (ctx.LESS_THAN() != null) {
                    return temStr1 + " < " + temStr2;
                } else if (ctx.GREATER_THAN() != null) {
                    return temStr1 + " > " + temStr2;
                } else if (ctx.EQUALS() != null) {
                    return temStr1 + " == " + temStr2;
                } else if (ctx.GT_EQ() != null) {
                    return temStr1 + " >= " + temStr2;
                } else if (ctx.LT_EQ() != null) {
                    return temStr1 + " <= " + temStr2;
                } else if (ctx.NOT_EQ_1() != null) {
                    return temStr1 + " <> " + temStr2;
                } else if (ctx.NOT_EQ_2() != null) {
                    return temStr1 + " != " + temStr2;
                } else if (ctx.IN() != null) {
                    if (ctx.optional != null && ctx.optional.getText() == "not") {
                        return temStr1 + " not in " + temStr2;
                    } else {
                        return temStr1 + " in " + temStr2;
                    }
                } else if (ctx.IS() != null) {
                    if (ctx.optional != null && ctx.optional.getText() == "not") {
                        return temStr1 + " is in " + temStr2;
                    } else {
                        return temStr1 + " in " + temStr2;
                    }
                }
            } else {
                // null
                return "";
            }
            return "";
        }
        @Override
        public String visitExpr(PythonParser.ExprContext ctx) {
            //    : AWAIT? atom trailer*
            //    | <assoc=right> expr op=POWER expr
            //    | op=(ADD | MINUS | NOT_OP) expr
            //    | expr op=(STAR | DIV | MOD | IDIV | AT) expr
            //    | expr op=(ADD | MINUS) expr
            //    | expr op=(LEFT_SHIFT | RIGHT_SHIFT) expr
            //    | expr op=AND_OP expr
            //    | expr op=XOR expr
            //    | expr op=OR_OP expr
            //    ;
            // judge the self var
            if (ctx.atom() != null) {
                String temStr = visit(ctx.atom());
                if (ctx.trailer() != null) {
                    for (PythonParser.TrailerContext temCtx : ctx.trailer()) {
                        temStr = temStr + visit(temCtx);
                    }
                }
                if (ctx.AWAIT() != null)
                    return "await " + temStr;
                else
                    return temStr;
            } else if (ctx.op != null) {
                if (ctx.expr().size() == 2){
                    // two sizes of arg
                    return visit(ctx.expr(0)) + ctx.op.getText() + visit(ctx.expr(1));
                } else {
                    // just one arg
                    return ctx.op.getText() + visit(ctx.expr(0));
                }
            } else {
                // no op
                return null;
            }
        }
        @Override
        public String visitAtom(PythonParser.AtomContext ctx) {
            //    : OPEN_PAREN (yield_expr | testlist_comp)? CLOSE_PAREN
            //    | OPEN_BRACKET testlist_comp? CLOSE_BRACKET
            //    | OPEN_BRACE dictorsetmaker? CLOSE_BRACE
            //    | REVERSE_QUOTE testlist COMMA? REVERSE_QUOTE
            //    | ELLIPSIS
            //    | name
            //    | PRINT
            //    | EXEC
            //    | MINUS? number
            //    | NONE
            //    | STRING+
            //    ;
            if (ctx.OPEN_PAREN() != null) {
                if (ctx.yield_expr() != null)
                    return '('+ visit(ctx.yield_expr()) +')';
                else if (ctx.testlist_comp() != null) {
                    return '('+ visit(ctx.testlist_comp()) +')';
                } else {
                    return "()";
                }
            } else if (ctx.OPEN_BRACKET() != null) {
                if (ctx.testlist_comp() != null) {
                    return '['+ visit(ctx.testlist_comp()) +']';
                }  else {
                    return "[]";
                }
            } else if (ctx.OPEN_BRACE() != null) {
                if (ctx.dictorsetmaker() != null) {
                    return '{' + visit(ctx.dictorsetmaker()) + '}';
                } else {
                    return "{}";
                }
            } else if (ctx.REVERSE_QUOTE() != null && ctx.REVERSE_QUOTE().size() != 0) {
                if (ctx.COMMA() != null) {
                    return "`" + visit(ctx.testlist()) + "," +"`";
                } else {
                    return "`" + visit(ctx.testlist()) +"`";
                }
            } else if (ctx.ELLIPSIS() != null) {
                return "...";
            } else if (ctx.name() != null) {
                return visit(ctx.name());
            } else if (ctx.PRINT() != null) {
                return "print";
            }else if (ctx.EXEC() != null) {
                return "exec";
            }else if (ctx.number() != null) {
                if (ctx.MINUS() != null) {
                    return "-" + visit(ctx.number());
                } else {
                    return visit(ctx.number());
                }
            } else if (ctx.NONE() != null) {
                return "None";
            } else if (ctx.STRING() != null) {
                String tem = "";
                for (int i = 0; i < ctx.STRING().size(); i++) {
                    tem = tem + ctx.STRING(i).getText();
                }
                return tem;
            }
            return "";
        }
        @Override
        public String visitDictorsetmaker(PythonParser.DictorsetmakerContext ctx) {
            //  : (test COLON test | POWER expr) (COMMA (test COLON test | POWER expr))* COMMA? // key_datum_list
            //    | test COLON test comp_for                                                      // dict_comprehension
            //    | testlist_comp
            //    ;
            if(ctx.testlist_comp() != null){
                return visit(ctx.testlist_comp());
            } else if (ctx.comp_for() != null) {
                return visit(ctx.test(0)) + "," + visit(ctx.test(1)) + visit(ctx.comp_for());
            } else {
                String temStr = "";
                for (PythonParser.TestContext temCtx : ctx.test()) {
                    temStr = temStr + visit(temCtx);
                }
                for (PythonParser.ExprContext temExpr : ctx.expr()) {
                    temStr = temStr + visit(temExpr);
                }
                return temStr;
            }
        }

        @Override
        public String visitTestlist(PythonParser.TestlistContext ctx) {
            //  : test (COMMA test)* COMMA?
            String temStr = "";
            for (PythonParser.TestContext temCtx : ctx.test()) {
                temStr = temStr + "," + visit(temCtx);
            }
            return temStr;
        }
        //=====================================================================//
        //                          PRIVATE METHODS                            //
        //=====================================================================//
        /**
         * Deal with the normal statements
         */
        private void visitStatement(ParserRuleContext ctx, String normalized) {
            ASNode statementNode = new ASNode(ASNode.Type.STATEMENT);
            statementNode.setCode(getOriginalCodeText(ctx));
            statementNode.setNormalizedCode(normalized);
            statementNode.setLineOfCode(ctx.getStart().getLine());
            AST.addVertex(statementNode);
            AST.addEdge(parentStack.peek(), statementNode);
        }
        @Override
        public String visitName(PythonParser.NameContext ctx) {
            //  : NAME
            //    | TRUE
            //    | FALSE
            //    ;
            if (ctx.NAME() != null)
                return ctx.NAME().getText();
            return ctx.getText();
        }
        @Override
        public String visitNumber(PythonParser.NumberContext ctx) {
            // number
            //    : integer
            //    | IMAG_NUMBER
            //    | FLOAT_NUMBER
            //    ;
            return ctx.getText();
        }
        @Override
        public String visitTest(PythonParser.TestContext ctx) {
            // : logical_test (IF logical_test ELSE test)?
            //    | LAMBDA varargslist? COLON test
            //    ;
            if (ctx.logical_test() != null) {
                // the first situation
                for (PythonParser.Logical_testContext temCtx : ctx.logical_test()) {
                    visit(temCtx);
                }
                if (ctx.test() != null) {
                    visit(ctx.test());
                }
            } else if (ctx.LAMBDA() != null) {
                // the second situation
                if (ctx.varargslist() != null)
                    return  "lambda " + visit(ctx.varargslist()) + " : " + visit(ctx.test());
            } else {
                // null
                return "";
            }
            return "";
        }

        /**
         * Get the original program text for the given parser-rule context.
         * This is required for preserving white-spaces.
         */
        private String getOriginalCodeText(ParserRuleContext ctx) {
            if (ctx == null)
                return "";
            int start = ctx.start.getStartIndex();
            int stop = ctx.stop.getStopIndex();
            Interval interval = new Interval(start, stop);
            return ctx.start.getInputStream().getText(interval);
        }
        // reset the function vars
        private void resetLocalVars() {
            vars.clear();
            varsCounter = 0;
        }

    }
}
